import torch
import PIL
import numpy
import base64
from torchvision import transforms
from io import BytesIO
import os


model_path = os.path.dirname(__file__) + '/recom.pth'

class Recom:
    def __init__(self):
        self.model = torch.jit.load(model_path).eval()
        self.trans = transforms.Compose([transforms.Resize((224, 224)),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.485, 0.456, 0.406],
                                                           [0.229, 0.224, 0.225])])
        self.model.cuda()
    
    def get_class(self, imStr: str):
        with torch.no_grad():
            pic_bin = base64.b64decode(imStr)
            pic = PIL.Image.open(BytesIO(pic_bin))
            
            img = self.trans(pic).unsqueeze(0).cuda()
            
            res = self.model(img)
            ind = res.argmax()
            return ind.item()
    
    @classmethod
    def get_label():
        label = [
        '人物',
        '动物',
        '植物',
        '自然风景',
        '人造物',
        '建筑'
        ]


if __name__ == '__main__':
    rf = open('test.jpg', 'rb').read()
    pic_str = base64.b64encode(rf)
    
    
    recom = Recom()
    
    res = recom.get_class(pic_str)
    print(res)
    